import time
import torch
import torch.nn.functional as F
import numpy as np

def Uncertainty_aware_Fusion(alpha_dict, classes, balance_term=False):
    """
    Uncertainty-aware Fusion (UF) to realize Collaborative Decision Making (CDM).
    alpha_dict: 包含各个分类器alpha参数的字典
    classes: 类别数量
    balance_term: 是否使用平衡项
    """
    def sub_uf_fast(alpha1, alpha2):
        # 使用向量化操作减少中间变量
        S1 = torch.sum(alpha1, dim=1, keepdim=True)
        S2 = torch.sum(alpha2, dim=1, keepdim=True)
        
        E1 = alpha1 - 1
        E2 = alpha2 - 1
        
        b1 = E1 / S1
        b2 = E2 / S2
        
        u1 = classes / S1
        u2 = classes / S2
        
        if balance_term:
            b_f = ((b1 + b2) * 0.5 + b1 * b2) * 0.5 + b1 * (1 - u1) + b2 * (1 - u2)
            u_f = u1 + u2 + u1 * u2
        else:
            b_f = b1 * b2 + b1 * (1 - u1) + b2 * (1 - u2)
            u_f = u1 * u2
        
        S_f = classes / u_f
        e_f = b_f * S_f
        alpha_f = e_f + 1
        
        return alpha_f
    
    # 批量处理所有分类器 - 修复：使用正确的参数名 alpha_dict
    alpha_list = list(alpha_dict.values())
    if len(alpha_list) == 1:
        return alpha_list[0]
    
    alpha_fuse = alpha_list[0]
    for i in range(1, len(alpha_list)):
        alpha_fuse = sub_uf_fast(alpha_fuse, alpha_list[i])
      
    return alpha_fuse